
import pandas as pd
import json
import re
from pathlib import Path


def extract_numbers_from_text(text):
    """从文本中提取数字信息"""
    if pd.isna(text) or not isinstance(text, str):
        return None
    
    # 提取人口数量的模式
    population_patterns = [
        r'(\d+(?:\.\d+)?)\s*万户',
        r'(\d+(?:\.\d+)?)\s*万口',
        r'(\d+(?:\.\d+)?)\s*万人',
        r'人口.*?(\d+(?:\.\d+)?)\s*万',
        r'户.*?(\d+(?:\.\d+)?)\s*万',
        r'(\d+)\s*户',
        r'(\d+)\s*口',
    ]
    
    # 提取年份的模式
    year_patterns = [
        r'(\d{3,4})\s*年',
        r'公元\s*(\d{3,4})',
        r'至元.*?(\d+)\s*年',
        r'大定.*?(\d+)\s*年',
        r'贞元.*?(\d+)\s*年',
    ]
    
    population = None
    area = None
    year = None
    
    # 提取人口
    for pattern in population_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                num = float(match.group(1))
                if '万' in pattern:
                    population = int(num * 10000)
                else:
                    population = int(num)
                break
            except:
                pass
    
    # 提取面积（按优先级从高到低，更精确的匹配在前）
    # 1. 平方千米/平方公里（直接使用）
    area_match = re.search(r'面积.*?(\d+(?:\.\d+)?)\s*(?:万\s*)?平方(?:千米|公里)', text)
    if not area_match:
        area_match = re.search(r'(\d+(?:\.\d+)?)\s*(?:万\s*)?平方(?:千米|公里)', text)
    
    if area_match:
        try:
            area_value = float(area_match.group(1))
            # 检查是否有"万"字
            if '万' in area_match.group(0):
                area = area_value * 10000  # 万平方千米
            else:
                area = area_value
        except:
            pass
    else:
        # 2. 耕地面积（亩）- 转换为平方千米
        area_match = re.search(r'耕地面积\s*(\d+)\s*亩', text)
        if area_match:
            try:
                area_mu = float(area_match.group(1))
                area = round(area_mu / 1500, 2)  # 1平方千米 = 1500亩
            except:
                pass
    
    # 提取年份
    for pattern in year_patterns:
        match = re.search(pattern, text)
        if match:
            try:
                year = int(match.group(1))
                # 处理年号转公元年的情况
                if year < 300:  # 可能是年号年份
                    continue
                break
            except:
                pass
    
    return {'population': population, 'area': area, 'year': year}


def extract_data_from_excel(file_path, sheet_name='总人口'):
    """从Excel文件中提取数据"""
    try:
        # 读取Excel文件
        df = pd.read_excel(file_path, sheet_name=sheet_name)
        
        # 如果第一行是标题行，使用它
        if df.columns[0] == '时期' or '时期' in str(df.iloc[0, 0]):
            # 检查是否需要重新设置列名
            if df.iloc[0, 0] == '时期':
                df.columns = df.iloc[0]
                df = df.iloc[1:].reset_index(drop=True)
        
        results = []
        
        for idx, row in df.iterrows():
            # 获取朝代
            category = row.iloc[0] if not pd.isna(row.iloc[0]) else None
            
            # 获取原文内容
            content = row.iloc[1] if len(row) > 1 else None
            
            if category and content:
                # 从文本中提取数字信息
                extracted = extract_numbers_from_text(str(content))
                
                # 创建数据记录
                record = {
                    'category': str(category),
                    'year': extracted['year'],
                    'population': extracted['population'],
                    'area': extracted['area']
                }
                
                # 只添加至少有一个数值字段的记录
                if any([record['year'], record['population'], record['area']]):
                    results.append(record)
        
        return results
    
    except Exception as e:
        print(f"读取文件时出错: {e}")
        return []


def main():
    # 输入文件路径
    input_file = './data/11人口 - 总数据和各朝代数据.xlsx'
    
    # 输出文件路径
    output_file = './outputs/population_data.json'
    
    print("开始提取数据...")
    
    # 提取第一个sheet的数据
    data = extract_data_from_excel(input_file, sheet_name='总人口')
    
    print(f"从'总人口' sheet提取了 {len(data)} 条记录")
    
    # 保存为JSON文件
    Path(output_file).parent.mkdir(parents=True, exist_ok=True)
    with open(output_file, 'w', encoding='utf-8') as f:
        json.dump(data, f, ensure_ascii=False, indent=2)

    print(f"数据已保存到: {output_file}")
    
    # 显示前几条数据作为示例
    print("\n前5条数据示例:")
    for i, record in enumerate(data[:5], 1):
        print(f"\n记录 {i}:")
        print(f"  朝代(category): {record['category']}")
        print(f"  年份(year): {record['year']}")
        print(f"  人口(population): {record['population']}")
        print(f"  面积(area): {record['area']} 平方千米" if record['area'] else f"  面积(area): {record['area']}")
    
    print(f"\n总共提取了 {len(data)} 条有效记录")
    print("注意：所有面积单位已统一为平方千米")
    
    return data


if __name__ == '__main__':
    main()
